from src import Preproduction as Pre

id, ch = Pre.load("results/hanzi_2_one_hot.data")
cnt = len(ch)
V = 10000

print("cnt = " + str(cnt))
print("V = " + str(V))

from src.EmbeddingGloVe import genX as genX
from src.EmbeddingGloVe import Glove as Glove

pathBase = 'dataset/chat_corpus/clean_chat_corpus/'
paths = [
    "chatterbot.tsv",
    "douban_single_turn.tsv",
    "ptt.tsv",
    "qingyun.tsv",
    "subtitle.tsv",
    "tieba.tsv",
    "weibo.tsv"
]
steps = [
    1, 
    300,
    40,
    10,
    274,
    232,
    443
]
fullPaths = [pathBase + path for path in paths]

import numpy as np

glove = Glove(V, 256)
glove.load("results/embedding120.model")

#i = glove.eval(id[c])

import random

# 从文档当中，采样batch个有监督文本串
# X=seqLen个汉字, Y=下一个汉字
# \t \n 视为END（id=0）
def Sample(rows, batch, seqLen, visualize=False):
    X = []
    Y = []
    
    for row in random.sample(rows, min(batch, len(rows))):
        for offset in range(1, len(row)):
            x = row[max(0, offset - seqLen) : offset]
            y = row[offset]
            if visualize:
                print(x + "> " + y)

            # 补充0
            x = [glove.eval(0).detach().numpy()] * max(0, seqLen - offset) + \
                [glove.eval(1 if i == '\t' or i == '\n' else id[i] if id[i] < V else 0).detach().numpy() for i in x]
            y = 1 if y == '\t' or y == '\n' else id[y] if id[y] < V else 0
            X.append(x)
            Y.append(y)

    return (np.array(X), np.array(Y))

# 样例
with open(fullPaths[0], "r") as f:
    rows = f.readlines()
    x, y = Sample(rows, 1, 10, True)
    print(x.shape, y.shape)

from src.GeneratorLSTM import LSTM

embDim = 256
hidDim = 256
seqLen = 20

lstm = LSTM(embDim, hidDim, seqLen, V)
#lstm.load("results/V3LSTM085.model")

# 训练
globalStep = 0
for i in range(1, 31):
    print("Epoch %d." % i)
    up = 0
    down = 0
    for j in range(len(paths)):
        print(paths[j])
        print("Total Steps: %d" % (steps[j] * 2))
        with open(fullPaths[j], "r") as f:
            rows = f.readlines()
            for step in range(steps[j] * 2):
                print("Epoch %d, Step %d." % (i, step))
                globalStep += 1
                X, Y = Sample(rows, 5, seqLen)
                loss = lstm.trainStep(X, Y)
                #print("Loss: " + str(loss.data))
                print('{{"metric": "loss", "value": {}, "step": {}}}'.format(loss.detach().numpy(), globalStep))
                print("------------------------------")    
            down += steps[j]
            up += steps[j] * loss.detach().numpy()
    print("Avg loss: %f" % (up / down))
    if i % 1 == 0:
        print("Saving...")
        lstm.save("results/V3LSTM%03d.model", i)
    print("==============================\n")

    